import os
import time
import csv
import subprocess
import re
import matplotlib.pyplot as plt
import pandas as pd

def get_gpu_stats():
    try:
        result = subprocess.run(['nvidia-smi', '--query-gpu=timestamp,temperature.gpu,fan.speed,utilization.gpu,utilization.memory,memory.total,memory.used,memory.free', '--format=csv,noheader,nounits'], capture_output=True, text=True, check=True)
        lines = result.stdout.strip().split('\n')
        stats = []

        for line in lines:
            values = line.split(', ')
            gpu_stat = {
                'timestamp': time.time(),
                'temperature': int(values[1]) if values[1] != '[N/A]' else None,
                'fan_speed': int(values[2]) if values[2] != '[N/A]' else None,
                'gpu_utilization': int(values[3]) if values[3] != '[N/A]' else None,
                'memory_utilization': int(values[4]) if values[4] != '[N/A]' else None,
                'total_memory': int(values[5]) / 1024.0 if values[5] != '[N/A]' else None,
                'used_memory': int(values[6]) / 1024.0 if values[6] != '[N/A]' else None,
                'free_memory': int(values[7]) / 1024.0 if values[7] != '[N/A]' else None
            }
            stats.append(gpu_stat)

        return stats

    except subprocess.CalledProcessError as e:
        print(f"Error running nvidia-smi: {e}")
        return []

def save_to_csv(stats, filename='gpu_stats.csv'):
    if not stats:
        return

    keys = stats[0].keys()
    with open(filename, mode='w', newline='') as file:
        writer = csv.DictWriter(file, fieldnames=keys)
        writer.writeheader()
        writer.writerows(stats)

def plot_stats(stats):
    if not stats:
        return

    df = pd.DataFrame(stats)
    df['timestamp'] = pd.to_datetime(df['timestamp'], unit='s')

    # Drop rows with any NaN values to avoid plotting issues
    df.dropna(inplace=True)

    fig, axs = plt.subplots(5, 1, figsize=(10, 20))

    axs[0].plot(df['timestamp'], df['temperature'])
    axs[0].set_title('GPU Temperature')
    axs[0].set_ylabel('Temperature (C)')

    axs[1].plot(df['timestamp'], df['fan_speed'])
    axs[1].set_title('GPU Fan Speed')
    axs[1].set_ylabel('Percentage (%)')

    axs[2].plot(df['timestamp'], df['gpu_utilization'])
    axs[2].set_title('GPU Utilization')
    axs[2].set_ylabel('Percentage (%)')

    axs[3].plot(df['timestamp'], df['memory_utilization'])
    axs[3].set_title('Memory Utilization')
    axs[3].set_ylabel('Percentage (%)')

    axs[4].plot(df['timestamp'], df['used_memory'])
    axs[4].set_title('Used Memory')
    axs[4].set_ylabel('Memory (GB)')

    for ax in axs:
        ax.set_xlabel('Time')
        ax.grid(True)

    plt.tight_layout()
    plt.show()

def plot_from_csv(filename='gpu_stats.csv'):
    try:
        df = pd.read_csv(filename)
        df['timestamp'] = pd.to_datetime(df['timestamp'], unit='s')

        # Drop rows with any NaN values to avoid plotting issues
        df.dropna(inplace=True)

        fig, axs = plt.subplots(5, 1, figsize=(10, 20))

        axs[0].plot(df['timestamp'], df['temperature'])
        axs[0].set_title('GPU Temperature')
        axs[0].set_ylabel('Temperature (C)')

        axs[1].plot(df['timestamp'], df['fan_speed'])
        axs[1].set_title('GPU Fan Speed')
        axs[1].set_ylabel('Percentage (%)')

        axs[2].plot(df['timestamp'], df['gpu_utilization'])
        axs[2].set_title('GPU Utilization')
        axs[2].set_ylabel('Percentage (%)')

        axs[3].plot(df['timestamp'], df['memory_utilization'])
        axs[3].set_title('Memory Utilization')
        axs[3].set_ylabel('Percentage (%)')

        axs[4].plot(df['timestamp'], df['used_memory'])
        axs[4].set_title('Used Memory')
        axs[4].set_ylabel('Memory (GB)')

        for ax in axs:
            ax.set_xlabel('Time')
            ax.grid(True)

        plt.tight_layout()
        plt.show()

    except FileNotFoundError:
        print(f"File {filename} not found.")
    except Exception as e:
        print(f"An error occurred while reading the CSV file: {e}")

if __name__ == "__main__":
    duration = 60  # Duration of the test in seconds
    interval = 1   # Interval between measurements in seconds

    all_stats = []
    start_time = time.time()

    while time.time() - start_time < duration:
        stats = get_gpu_stats()
        all_stats.extend(stats)
        time.sleep(interval)

    save_to_csv(all_stats)
    # plot_stats(all_stats)

    # Uncomment the following line to plot from an existing CSV file
    # plot_from_csv('gpu_stats.csv')



